{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Llama3.1-8B-Instruct Reinforcement Learning Demo\n",
    "\n",
    "This notebook demonstrates training on Llama3.1-8B-Instruct model with either GRPO (Group Relative Policy Optimization) or GSPO (Group Sequence Policy Optimization).\n",
    "\n",
    "## What is GRPO/GSPO?\n",
    "\n",
    "GRPO/GSPO is an RL algorithm that enhances reasoning abilities of LLMs by:\n",
    "1. Generating multiple responses for each prompt\n",
    "2. Evaluating responses using reward models  \n",
    "3. Calculating relative advantages to update the policy\n",
    "\n",
    "The difference is in the loss function - either it's optimizing each token (GRPO) or the whole sequence(GSPO).\n",
    "\n",
    "## Hardware Requirements\n",
    "\n",
    "- Single host TPUVM (v6e-8/v5p-8)\n",
    "- Sufficient memory for Llama3.1-8B model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get Your Hugging Face Token\n",
    "\n",
    "To access model checkpoint from the Hugging Face Hub, you need to authenticate with a personal access token.\n",
    "\n",
    "**Follow these steps to get your token:**\n",
    "\n",
    "1.  **Navigate to the Access Tokens page** in your Hugging Face account settings. You can go there directly by visiting this URL:\n",
    "    *   [https://huggingface.co/settings/tokens](https://huggingface.co/settings/tokens)\n",
    "\n",
    "2.  **Create a new token** by clicking the **\"+ Create new token\"** button.\n",
    "\n",
    "3.  **Give your token a name** and assign it a **`read` role**. The `read` role is sufficient for downloading models.\n",
    "\n",
    "4.  **Copy the generated token**. You will need to paste it in the next step.\n",
    "\n",
    "**Follow these steps to store your token:**\n",
    "\n",
    "Just put your token in the line below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HF_TOKEN = \"\" # Set HF_TOKEN environment variable\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "Install dependencies and set up the environment:\n",
    "https://maxtext.readthedocs.io/en/latest/tutorials/posttraining/rl.html#from-github"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configuration\n",
    "\n",
    "Set up the training parameters. We use a single host TPU. Defaults are hardcoded for Llama3.1-8B:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if you have cloned the maxtext repo, you should set the path to the maxtext/src folder\n",
    "# otherwise, you can just run the cell below\n",
    "!cd ~/maxtext/src/  #  This is the path to the maxtext/src folder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Choose the loss algorithm between GSPO or GRPO\n",
    "LOSS_ALGO=\"grpo\" #  or \"gspo-token\" if you want to use GSPO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "from pathlib import Path\n",
    "import MaxText\n",
    "from huggingface_hub import login\n",
    "import jax\n",
    "\n",
    "# Set up paths (adjust if needed)\n",
    "MAXTEXT_REPO_ROOT = os.path.dirname(MaxText.__file__)\n",
    "RUN_NAME=\"grpo_test\"\n",
    "# Hardcoded defaults for Llama3.1-8B\n",
    "MODEL_NAME = \"llama3.1-8b\"\n",
    "HF_REPO_ID = \"meta-llama/Llama-3.1-8B-Instruct\"\n",
    "CHAT_TEMPLATE_PATH = f\"{MAXTEXT_REPO_ROOT}/examples/chat_templates/gsm8k_rl.json\"\n",
    "\n",
    "# Required: Set these before running\n",
    "MODEL_CHECKPOINT_PATH = \"\"  # Update this!\n",
    "if not MODEL_CHECKPOINT_PATH:\n",
    "    raise RuntimeError(\"MODEL_CHECKPOINT_PATH is not set\")\n",
    "    \n",
    "OUTPUT_DIRECTORY = \"\"  # Update this!\n",
    "if not OUTPUT_DIRECTORY:\n",
    "    raise RuntimeError(\"OUTPUT_DIRECTORY is not set\")\n",
    "    \n",
    "os.environ[\"HF_TOKEN\"] = HF_TOKEN\n",
    "\n",
    "if HF_TOKEN:\n",
    "    login(token=HF_TOKEN)\n",
    "    print(\"Authenticated with Hugging Face\")\n",
    "else:\n",
    "    print(\"Authentication failed: Hugging Face token not set\")\n",
    "\n",
    "# Optional: Override training parameters\n",
    "LEARNING_RATE = 3e-6\n",
    "NUM_GENERATIONS = 2\n",
    "GRPO_BETA = 0.08\n",
    "GRPO_EPSILON = 0.2\n",
    "CHIPS_PER_VM = 1\n",
    "\n",
    "print(f\"📁 MaxText Home: {MAXTEXT_REPO_ROOT}\")\n",
    "print(f\"🤖 Model: {MODEL_NAME}\")\n",
    "print(f\"📦 Checkpoint: {MODEL_CHECKPOINT_PATH}\")\n",
    "print(f\"💾 Output: {OUTPUT_DIRECTORY}\")\n",
    "print(f\"🔑 HF Token: {'✅ Set' if HF_TOKEN else '❌ Missing - set HF_TOKEN env var'}\")\n",
    "print(f\"📊 Steps: {STEPS}\")\n",
    "print(f\"Loss Algorithm : {LOSS_ALGO}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add MaxText to Python path\n",
    "maxtext_path = Path(MAXTEXT_REPO_ROOT) \n",
    "sys.path.insert(0, str(maxtext_path))\n",
    "\n",
    "from MaxText import pyconfig, max_utils\n",
    "from MaxText.rl.train_rl import rl_train, setup_configs_and_devices\n",
    "\n",
    "# Initialize JAX \n",
    "os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"0\"\n",
    "os.environ[\"SKIP_JAX_PRECOMPILE\"] = \"1\"  # Faster startup for vLLM\n",
    "\n",
    "print(\"✅ Successfully imported modules\")\n",
    "print(f\"📁 MaxText path: {maxtext_path}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build configuration for GRPO training\n",
    "config_file = os.path.join(MAXTEXT_REPO_ROOT, \"configs/rl.yml\")\n",
    "\n",
    "# Verify chat template exists\n",
    "if not os.path.exists(CHAT_TEMPLATE_PATH)):\n",
    "    raise FileNotFoundError(f\"Chat template not found: {CHAT_TEMPLATE_PATH}\")\n",
    "\n",
    "# Build argv list for pyconfig.initialize()\n",
    "config_argv = [\n",
    "    \"\",  # argv[0] placeholder\n",
    "    config_file,\n",
    "    f\"model_name={MODEL_NAME}\",\n",
    "    f\"tokenizer_path={HF_REPO_ID}\",\n",
    "    f\"run_name={RUN_NAME}\",\n",
    "    f\"chat_template_path={CHAT_TEMPLATE_PATH}\",\n",
    "    f\"load_parameters_path={MODEL_CHECKPOINT_PATH}\",\n",
    "    f\"base_output_directory={OUTPUT_DIRECTORY}\",\n",
    "    f\"hf_access_token={HF_TOKEN}\",\n",
    "    f\"learning_rate={LEARNING_RATE}\",\n",
    "    f\"num_generations={NUM_GENERATIONS}\",\n",
    "    f\"grpo_beta={GRPO_BETA}\",\n",
    "    f\"grpo_epsilon={GRPO_EPSILON}\",\n",
    "    f\"chips_per_vm={CHIPS_PER_VM}\",\n",
    "    f\"loss_algo={LOSS_ALGO}\",\n",
    "    \"use_pathways=False\"\n",
    "]\n",
    "\n",
    "# Initialize configuration\n",
    "print(f\"🔧 Initializing configuration from: {config_file}\")\n",
    "config = pyconfig.initialize(config_argv)\n",
    "\n",
    "print(\"\\n✅ Configuration initialized successfully\")\n",
    "print(f\"📊 Training steps: {config.steps}\")\n",
    "print(f\"📁 Output directory: {config.base_output_directory}\")\n",
    "print(f\"🤖 Model: {config.model_name}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Execute GRPO/GSPO training\n",
    "print(\"\\n\" + \"=\"*80)\n",
    "print(\"🚀 Starting Training...\")\n",
    "print(\"=\"*80)\n",
    "print(1)\n",
    "try:\n",
    "    # Call the rl_train function (it handles everything internally)\n",
    "    rl_train(config)\n",
    "    \n",
    "    print(\"\\n\" + \"=\"*80)\n",
    "    print(\"✅ Training Completed Successfully!\")\n",
    "    print(\"=\"*80)\n",
    "    print(f\"📁 Checkpoints saved to: {config.checkpoint_dir}\")\n",
    "    print(f\"📊 TensorBoard logs: {config.tensorboard_dir}\")\n",
    "    print(f\"🎯 Model ready for inference!\")\n",
    "    \n",
    "except Exception as e:\n",
    "    print(\"\\n\" + \"=\"*80)\n",
    "    print(\"❌Training Failed!\")\n",
    "    print(\"=\"*80)\n",
    "    print(f\"Error: {str(e)}\")\n",
    "    import traceback\n",
    "    traceback.print_exc()\n",
    "    print(\"\\n💡 Common issues:\")\n",
    "    print(\"  - Check that MODEL_CHECKPOINT_PATH points to a valid checkpoint\")\n",
    "    print(\"  - Ensure HF_TOKEN environment variable is set\")\n",
    "    print(\"  - Verify OUTPUT_DIRECTORY is writable\")\n",
    "    print(\"  - Check hardware requirements (TPU/GPU availability)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 📚 Learn More\n",
    "\n",
    "- **CLI Usage**: https://maxtext.readthedocs.io/en/latest/tutorials/rl.html#run-grpo\n",
    "- **Configuration**: See `src/MaxText/configs/rl.yml` for all available options\n",
    "- **Documentation**: Check `src/MaxText/rl/train_rl.py` for the `rl_train` function implementation"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
